#include "contextchain.h"
#include <cassert>

/*
 * @brief 将数据数组中符合 `from` 模数下的负一元素转换为 `to` 模数下的负一元素。
 *
 * @param data 要处理的数据数组。
 * @param len  数据数组的长度。
 * @param from 原模数，用于计算负一。
 * @param to 目标模数，转换后的负一将对应于该模数。
 */
void ContextChain::ConvertNegOnes(uint64_t *data, size_t len, const uint64_t &from, const uint64_t &to)
{
    uint64_t fromNegOne = NegateUintMod(1, from);
    uint64_t toNegOne = NegateUintMod(1, to);
    for (int i = 0; i < len; i++) {
        if (data[i] == fromNegOne) {
            data[i] = toNegOne;
        }
    }
}
/*
 * @brief 构造函数，根据最基本的参数创建参数列表
 *
 * @param base_context 最基本的参数，明文模数为p
 * @param e 低位提取的次数，也是列表参数中明文模数的最高次数
 */
ContextChain::ContextChain(const Context &baseContext, size_t e)
{
    Context *parms = new Context(baseContext.logN, baseContext.qi, baseContext.pi, baseContext.t, baseContext.sigma);
    uint64_t p = parms->t;
    contexts.push_back(*parms);
    evaluators.emplace_back(std::make_unique<Evaluator>(parms));
    for (size_t i = 1; i < e; ++i) {
        Context *new_parms = new Context(parms->logN, parms->qi, parms->pi, Exponentiate(p, i + 1), parms->sigma);
        contexts.push_back(*new_parms); // 将新的 Context 对象存储到 contexts
        evaluators.emplace_back(std::make_unique<Evaluator>(new_parms));
    }
    assert(evaluators.size() == contexts.size());
}

size_t ContextChain::PolyModulusDegree() const
{
    return 1 << contexts.begin()->logN;
}
const Context &ContextChain::GetContext(size_t i) const
{
    return contexts[i];
}
Evaluator &ContextChain::GetEvaluator(size_t i) const
{
    return *evaluators[i];
}

size_t ContextChain::Size() const
{
    return contexts.size();
}

Context &ContextChain::BaseContext() const
{
    return const_cast<Context &>(*contexts.begin());
}

Context &ContextChain::TargetContext() const
{
    return const_cast<Context &>(contexts.back());
}

Evaluator &ContextChain::TargetEvaluator() const
{
    return GetEvaluator(Size() - 1);
}
/*
 * @brief 将密钥source转化成另一个密钥destination，密钥值不变，只有模有区别
 *
 * @param source 原始密钥
 * @param destination 转化后的密钥
 */
void ContextChain::ConvertSk(const SecretKey &source, SecretKey &destination) const
{
    const size_t polyModulusDegree = this->PolyModulusDegree();
    destination = source;
    Poly aa(source.value.Q->Coeffs[0].size(), 1);
    CopyValuesLvl(0, source.value.Q, &aa);
    TargetEvaluator().ringQ->InvNTTLvl(0, &aa, &aa);
    for (int i = 0; i < TargetContext().QCount(); i++) {
        destination.value.Q->Coeffs[i] = aa.Coeffs[0];
        ConvertNegOnes(&destination.value.Q->Coeffs[i][0], polyModulusDegree, TargetContext().qi[0],
            TargetContext().qi[i]);
    }
    TargetEvaluator().ringQ->NTT(destination.value.Q, destination.value.Q);
}
// 将密钥s的mod Q0的向量转化成明文mod p^e
void ContextChain::ConvertSkPlain(const SecretKey &source, Plaintext &destination) const
{
    const size_t polyModulusDegree = this->PolyModulusDegree();
    Plaintext skPlain(source.value.Q->Coeffs[0], TargetContext().MaxLevel());
    TargetEvaluator().ringQ->InvNTTLvl(0, skPlain.value, skPlain.value);
    destination.value->Coeffs[0] = skPlain.value->Coeffs[0];
    ConvertNegOnes(&destination.value->Coeffs[0][0], polyModulusDegree, TargetContext().ringQ->Modulus[0],
        TargetContext().T());
}

// 将密文的两个多项式（mod Q）分别转换到两个明文 （mod p^e）上
void ContextChain::ConvertCiphertextPlain(const Ciphertext &source, Plaintext *destination) const
{
    if (source.value.size() != 2) {
        throw std::invalid_argument("Can only convert size 2 ciphertexts to plaintexts");
    }
    for (int i = 0; i < 2; i++) {
        TargetEvaluator().lightEncoder->scaler->DivByQOverTRounded(source.value[i], destination[i].value);
    }
}
